{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And now - Week 3 Day 3\n",
    "\n",
    "## AutoGen Core\n",
    "\n",
    "Something a little different.\n",
    "\n",
    "This is agnostic to the underlying Agent framework\n",
    "\n",
    "You can use AutoGen AgentChat, or you can use something else; it's an Agent interaction framework.\n",
    "\n",
    "From that point of view, it's positioned similarly to LangGraph.\n",
    "\n",
    "### The fundamental principle\n",
    "\n",
    "Autogen Core decouples an agent's logic from how messages are delivered.  \n",
    "The framework provides a communication infrastructure, along with agent lifecycle, and the agents are responsible for their own work.\n",
    "\n",
    "The communication infrastructure is called a Runtime.\n",
    "\n",
    "There are 2 types: **Standalone** and **Distributed**.\n",
    "\n",
    "Today we will use a standalone runtime: the **SingleThreadedAgentRuntime**, a local embedded agent runtime implementation.\n",
    "\n",
    "Tomorrow we'll briefly look at a Distributed runtime.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
    "from autogen_core import SingleThreadedAgentRuntime\n",
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_agentchat.messages import TextMessage\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(override=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First we define our Message object\n",
    "\n",
    "Whatever structure we want for messages in our Agent framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's have a simple one!\n",
    "\n",
    "@dataclass\n",
    "class Message:\n",
    "    content: str\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now we define our Agent\n",
    "\n",
    "A subclass of RoutedAgent.\n",
    "\n",
    "Every Agent has an **Agent ID** which has 2 components:  \n",
    "`agent.id.type` describes the kind of agent it is  \n",
    "`agent.id.key` gives it its unique identifier\n",
    "\n",
    "Any method with the `@message_handler` decorated will have the opportunity to receive messages.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleAgent(RoutedAgent):\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__(\"Simple\")\n",
    "\n",
    "    @message_handler\n",
    "    async def on_my_message(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        return Message(content=f\"This is {self.id.type}-{self.id.key}. You said '{message.content}' and I disagree.\")\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OK let's create a Standalone runtime and register our agent type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "runtime = SingleThreadedAgentRuntime()\n",
    "await SimpleAgent.register(runtime, \"simple_agent\", lambda: SimpleAgent())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alright! Let's start a runtime and send a message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "runtime.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_id = AgentId(\"simple_agent\", \"default\")\n",
    "response = await runtime.send_message(Message(\"Well hi there!\"), agent_id)\n",
    "print(\">>>\", response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "await runtime.stop()\n",
    "await runtime.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OK Now let's do something more interesting\n",
    "\n",
    "We'll use an AgentChat Assistant!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class MyLLMAgent(RoutedAgent):\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__(\"LLMAgent\")\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n",
    "        self._delegate = AssistantAgent(\"LLMAgent\", model_client=model_client)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        print(f\"{self.id.type} received message: {message.content}\")\n",
    "        text_message = TextMessage(content=message.content, source=\"user\")\n",
    "        response = await self._delegate.on_messages([text_message], ctx.cancellation_token)\n",
    "        reply = response.chat_message.content\n",
    "        print(f\"{self.id.type} responded: {reply}\")\n",
    "        return Message(content=reply)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_core import SingleThreadedAgentRuntime\n",
    "\n",
    "runtime = SingleThreadedAgentRuntime()\n",
    "await SimpleAgent.register(runtime, \"simple_agent\", lambda: SimpleAgent())\n",
    "await MyLLMAgent.register(runtime, \"LLMAgent\", lambda: MyLLMAgent())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runtime.start()  # Start processing messages in the background.\n",
    "response = await runtime.send_message(Message(\"Hi there!\"), AgentId(\"LLMAgent\", \"default\"))\n",
    "print(\">>>\", response.content)\n",
    "response =  await runtime.send_message(Message(response.content), AgentId(\"simple_agent\", \"default\"))\n",
    "print(\">>>\", response.content)\n",
    "response = await runtime.send_message(Message(response.content), AgentId(\"LLMAgent\", \"default\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "await runtime.stop()\n",
    "await runtime.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OK now let's show this at work - let's have 3 agents interact!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_ext.models.ollama import OllamaChatCompletionClient\n",
    "\n",
    "\n",
    "class Player1Agent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\", temperature=1.0)\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        text_message = TextMessage(content=message.content, source=\"user\")\n",
    "        response = await self._delegate.on_messages([text_message], ctx.cancellation_token)\n",
    "        return Message(content=response.chat_message.content)\n",
    "    \n",
    "class Player2Agent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OllamaChatCompletionClient(model=\"llama3.2\", temperature=1.0)\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        text_message = TextMessage(content=message.content, source=\"user\")\n",
    "        response = await self._delegate.on_messages([text_message], ctx.cancellation_token)\n",
    "        return Message(content=response.chat_message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "JUDGE = \"You are judging a game of rock, paper, scissors. The players have made these choices:\\n\"\n",
    "\n",
    "class RockPaperScissorsAgent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\", temperature=1.0)\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        instruction = \"You are playing rock, paper, scissors. Respond only with the one word, one of the following: rock, paper, or scissors.\"\n",
    "        message = Message(content=instruction)\n",
    "        inner_1 = AgentId(\"player1\", \"default\")\n",
    "        inner_2 = AgentId(\"player2\", \"default\")\n",
    "        response1 = await self.send_message(message, inner_1)\n",
    "        response2 = await self.send_message(message, inner_2)\n",
    "        result = f\"Player 1: {response1.content}\\nPlayer 2: {response2.content}\\n\"\n",
    "        judgement = f\"{JUDGE}{result}Who wins?\"\n",
    "        message = TextMessage(content=judgement, source=\"user\")\n",
    "        response = await self._delegate.on_messages([message], ctx.cancellation_token)\n",
    "        return Message(content=result + response.chat_message.content)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "runtime = SingleThreadedAgentRuntime()\n",
    "await Player1Agent.register(runtime, \"player1\", lambda: Player1Agent(\"player1\"))\n",
    "await Player2Agent.register(runtime, \"player2\", lambda: Player2Agent(\"player2\"))\n",
    "await RockPaperScissorsAgent.register(runtime, \"rock_paper_scissors\", lambda: RockPaperScissorsAgent(\"rock_paper_scissors\"))\n",
    "runtime.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_id = AgentId(\"rock_paper_scissors\", \"default\")\n",
    "message = Message(content=\"go\")\n",
    "response = await runtime.send_message(message, agent_id)\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "await runtime.stop()\n",
    "await runtime.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
